"""
python fujian1_arima_gragh.py
"""


import os
import json
import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm

# 定义路径
original_data_dir = 'fujian/fujian1/json_output'
predicted_data_dir = 'fujian/fujian1/predict_arima'
output_dir = 'fujian/fujian1/predict_arima/graph'

# 创建输出目录
os.makedirs(output_dir, exist_ok=True)

# 获取原始数据和预测数据文件列表
original_files = [f for f in os.listdir(original_data_dir) if f.endswith('.json')]
predicted_files = [f for f in os.listdir(predicted_data_dir) if f.endswith('.json')]

# 创建用于存储数据的字典
data_dict = {}

# 处理原始数据
for file in original_files:
    with open(os.path.join(original_data_dir, file), 'r') as f:
        data = json.load(f)
        seller_product_wh = file.split('.')[0]  # 从文件名获取组名
        dates = [entry['date'] for entry in data]  # 获取日期
        qty = [entry['qty'] for entry in data]  # 获取数量
        data_dict[seller_product_wh] = {
            'original_dates': pd.to_datetime(dates),
            'original_qty': qty
        }

# 处理预测数据
for file in predicted_files:
    with open(os.path.join(predicted_data_dir, file), 'r') as f:
        data = json.load(f)
        seller_product_wh = file.replace('predicted_processed_', '').replace('.json', '')  # 从文件名获取组名
        dates = [entry['date'] for entry in data]  # 获取日期
        qty = [entry['qty'] for entry in data]  # 获取数量
        if seller_product_wh in data_dict:
            data_dict[seller_product_wh]['predicted_dates'] = pd.to_datetime(dates)
            data_dict[seller_product_wh]['predicted_qty'] = qty

# 绘制图形并保存
for seller_product_wh, values in tqdm(data_dict.items(), desc='Processing groups'):
    if 'predicted_dates' in values:  # 确保有预测数据
        plt.figure(figsize=(12, 6))
        plt.plot(values['original_dates'], values['original_qty'], label='Original Data', color='blue')
        plt.plot(values['predicted_dates'], values['predicted_qty'], label='Predicted Data', color='orange')
        plt.title(f'Time Series for {seller_product_wh}')
        plt.xlabel('Date')
        plt.ylabel('Quantity')
        plt.legend()
        plt.xticks(rotation=45)
        plt.tight_layout()

        # 保存图形
        plt.savefig(os.path.join(output_dir, f'{seller_product_wh}.png'))
        plt.close()

print("All graphs have been saved successfully!")
